#include <ComputeShaderHistogram/ComputeShaderHistogram.h>

#include <Core/Graphics/Geometry.h>
#include <Core/Input/DeviceTypes/MouseKeyboard.h>
#include <Core/Input/InputManager.h>
#include <Core/System/Window.h>
#include <Core/System/WindowManager.h>
#include <Foundation/Time/Clock.h>
#include <Foundation/Utilities/Stats.h>
#include <GameEngine/GameApplication/WindowOutputTarget.h>
#include <RendererCore/RenderContext/RenderContext.h>
#include <RendererCore/ShaderCompiler/ShaderManager.h>
#include <RendererFoundation/Device/SwapChain.h>
#include <RendererFoundation/Resources/Texture.h>

#include "../../../Data/Samples/ComputeShaderHistogram/Shaders/ClearHistogramParams.h"

static ezUInt32 g_uiComputeThreadGroupSize = 32;

ezComputeShaderHistogramApp::ezComputeShaderHistogramApp()
  : ezGameApplication("ComputeShaderHistogram", "Data/Samples/ComputeShaderHistogram")

{
}

ezComputeShaderHistogramApp::~ezComputeShaderHistogramApp() = default;

void ezComputeShaderHistogramApp::Run()
{
  ezWindowManager::GetSingleton()->Update();

  ezClock::GetGlobalClock()->Update();
  Run_InputUpdate();

#if EZ_ENABLED(EZ_SUPPORTS_DIRECTORY_WATCHER)
  m_bStuffChanged = false;
  m_pDirectoryWatcher->EnumerateChanges(ezMakeDelegate(&ezComputeShaderHistogramApp::OnFileChanged, this));
  if (m_bStuffChanged)
  {
    ezResourceManager::ReloadAllResources(false);
  }
#endif

  const ezUInt32 uiWindowWidth = m_pWindow->GetClientAreaSize().width;
  const ezUInt32 uiWindowHeight = m_pWindow->GetClientAreaSize().height;

  // do the rendering
  {
    auto device = ezGALDevice::GetDefaultDevice();

    const ezGALSwapChain* pPrimarySwapChain = device->GetSwapChain(m_hSwapChain);
    ezGALRenderTargetViewHandle hBackbufferRTV = device->GetDefaultRenderTargetView(pPrimarySwapChain->GetBackBufferTexture());

    // Before starting to render in a frame call this function
    device->EnqueueFrameSwapChain(m_hSwapChain);
    device->BeginFrame();

    ezGALCommandEncoder* pCommandEncoder = device->BeginCommands("ezComputeShaderHistogram");

    ezRenderContext& renderContext = *ezRenderContext::GetDefaultInstance();

    // Constant buffer update
    {
      auto& globalConstants = renderContext.WriteGlobalConstants();
      ezMemoryUtils::ZeroFill(&globalConstants, 1);

      globalConstants.ViewportSize = ezVec4((float)uiWindowWidth, (float)uiWindowHeight, 1.0f / (float)uiWindowWidth, 1.0f / (float)uiWindowHeight);
      // Wrap around to prevent floating point issues. Wrap around is dividable by all whole numbers up to 11.
      globalConstants.GlobalTime = (float)ezMath::Mod(ezClock::GetGlobalClock()->GetAccumulatedTime().GetSeconds(), 20790.0);
      globalConstants.WorldTime = globalConstants.GlobalTime;
    }

    ezRectFloat viewport(0.0f, 0.0f, (float)uiWindowWidth, (float)uiWindowHeight);

    // Draw background.
    {
      ezGALRenderingSetup renderingSetup;
      renderingSetup.SetColorTarget(0, m_hScreenRTV);
      renderContext.BeginRendering(renderingSetup, viewport, "Background");

      renderContext.BindShader(m_hRenderScreenShader);
      renderContext.BindNullMeshBuffer(ezGALPrimitiveTopology::Triangles, 1); // Vertices are generated by shader.
      renderContext.DrawMeshBuffer().IgnoreResult();


      renderContext.EndRendering();
    }

    // Render screen texture to backbuffer
    {
      ezGALRenderingSetup renderingSetup;
      renderingSetup.SetColorTarget(0, hBackbufferRTV);
      renderContext.BeginRendering(renderingSetup, viewport, "Blit");
      ezBindGroupBuilder& bindGroupSample = ezRenderContext::GetDefaultInstance()->GetBindGroup();
      bindGroupSample.BindTexture("Input", m_hScreenTexture);
      renderContext.BindShader(m_hDisplayScreenShader);
      renderContext.BindNullMeshBuffer(ezGALPrimitiveTopology::Triangles, 1);
      renderContext.DrawMeshBuffer().IgnoreResult();
      renderContext.EndRendering();
    }

    // Compute histogram.
    {
      renderContext.BeginCompute("ComputeHistogram");

      ezBindGroupBuilder& bindGroupSample = ezRenderContext::GetDefaultInstance()->GetBindGroup();
      // Reset first.
      renderContext.BindShader(m_hClearHistogramShader);
      ClearHistogramParams params;
      params.Height = 3;
      params.Width = 256;
      renderContext.SetPushConstants("ClearHistogramParams", params);
      bindGroupSample.BindTexture("HistogramOutput", m_hHistogramTexture);
      renderContext.Dispatch(256 / 16, 1, 1).IgnoreResult();

      // Compute histogram
      renderContext.BindShader(m_hComputeHistogramShader);
      bindGroupSample.BindTexture("ScreenTexture", m_hScreenTexture);
      bindGroupSample.BindTexture("HistogramOutput", m_hHistogramTexture);
      renderContext.Dispatch(uiWindowWidth / g_uiComputeThreadGroupSize + (uiWindowWidth % g_uiComputeThreadGroupSize != 0 ? 1 : 0), uiWindowHeight / g_uiComputeThreadGroupSize + (uiWindowHeight % g_uiComputeThreadGroupSize != 0 ? 1 : 0)).IgnoreResult();

      renderContext.EndCompute();
    }

    // Draw histogram.
    {
      ezGALRenderingSetup renderingSetup;
      renderingSetup.SetColorTarget(0, hBackbufferRTV);
      renderContext.BeginRendering(renderingSetup, viewport, "DrawHistogram");

      renderContext.BindShader(m_hDisplayHistogramShader);
      renderContext.BindMeshBuffer(m_hHistogramQuadMeshBuffer);
      ezBindGroupBuilder& bindGroupSample = renderContext.GetBindGroup();
      bindGroupSample.BindTexture("HistogramTexture", m_hHistogramTexture);
      renderContext.DrawMeshBuffer().IgnoreResult();

      renderContext.EndRendering();
    }

    device->EndCommands(pCommandEncoder);

    device->EndFrame();
  }

  // needs to be called once per frame
  ezResourceManager::PerFrameUpdate();

  // tell the task system to finish its work for this frame
  // this has to be done at the very end, so that the task system will only use up the time that is left in this frame for
  // uploading GPU data etc.
  ezTaskSystem::FinishFrameTasks();
}

void ezComputeShaderHistogramApp::AfterCoreSystemsStartup()
{
  SUPER::AfterCoreSystemsStartup();
#if EZ_ENABLED(EZ_SUPPORTS_DIRECTORY_WATCHER)
  m_pDirectoryWatcher = EZ_DEFAULT_NEW(ezDirectoryWatcher);
  EZ_VERIFY(m_pDirectoryWatcher->OpenDirectory(FindProjectDirectory(), ezDirectoryWatcher::Watch::Writes | ezDirectoryWatcher::Watch::Subdirectories).Succeeded(), "Failed to watch project directory");
#endif
  auto device = ezGALDevice::GetDefaultDevice();

  if (auto pInput = ezInputManager::GetInputDeviceOfType<ezInputDeviceMouseKeyboard>())
  {
    pInput->SetShowMouseCursor(true);
    pInput->SetClipMouseCursor(ezMouseCursorClipMode::NoClip);
  }

  // Retrieve window and swapchain handle for rendering
  {
    auto pWinMan = ezWindowManager::GetSingleton();

    ezHybridArray<ezRegisteredWndHandle, 8> windowIds;
    pWinMan->GetRegistered(windowIds);

    for (auto id : windowIds)
    {
      m_pWindow = pWinMan->GetWindow(id);

      // Retrieve only the first output target
      if (auto pOutput = pWinMan->GetOutputTarget(id))
      {
        m_hSwapChain = static_cast<ezWindowOutputTargetGAL*>(pOutput)->m_hSwapChain;
        break;
      }
    }

    EZ_ASSERT_DEV(m_pWindow != nullptr, "Failed to retrieve active window. No window plugins have been registered.");
    EZ_ASSERT_DEV(!m_hSwapChain.IsInvalidated(), "Failed to retrieve active window output target.");
  }

  // Create textures and texture view for screen content (can't use back-buffer as shader resource view)
  {
    ezGALTextureCreationDescription texDesc;
    texDesc.m_uiWidth = m_pWindow->GetClientAreaSize().width;
    texDesc.m_uiHeight = m_pWindow->GetClientAreaSize().height;
    texDesc.m_Format = ezGALResourceFormat::RGBAUByteNormalized; // ezGALResourceFormat::RGBAUByteNormalizedsRGB;
    texDesc.m_bAllowRenderTargetView = true;
    texDesc.m_bAllowShaderResourceView = true;

    m_hScreenTexture = device->CreateTexture(texDesc);
    m_hScreenRTV = device->GetDefaultRenderTargetView(m_hScreenTexture);
  }

  // Create texture for histogram data.
  {
    ezGALTextureCreationDescription texDesc;
    texDesc.m_uiWidth = 256;
    texDesc.m_uiHeight = 3; // R, G, B
    texDesc.m_uiMipLevelCount = 1;
    texDesc.m_Format = ezGALResourceFormat::RUInt;
    texDesc.m_bAllowRenderTargetView = false;
    texDesc.m_bAllowShaderResourceView = true;
    texDesc.m_bAllowUAV = true;
    texDesc.m_ResourceAccess.m_bImmutable = false;

    m_hHistogramTexture = device->CreateTexture(texDesc);
  }

  // Setup Shaders and Materials
  {
    m_hRenderScreenShader = ezResourceManager::LoadResource<ezShaderResource>("Shaders/RenderScreen.ezShader");
    m_hDisplayScreenShader = ezResourceManager::LoadResource<ezShaderResource>("Shaders/DisplayScreen.ezShader");
    m_hClearHistogramShader = ezResourceManager::LoadResource<ezShaderResource>("Shaders/ClearHistogram.ezShader");
    m_hComputeHistogramShader = ezResourceManager::LoadResource<ezShaderResource>("Shaders/ComputeHistogram.ezShader");
    m_hDisplayHistogramShader = ezResourceManager::LoadResource<ezShaderResource>("Shaders/DisplayHistogram.ezShader");
  }

  // Geometry.
  CreateHistogramQuad();
}

void ezComputeShaderHistogramApp::BeforeHighLevelSystemsShutdown()
{
  auto device = ezGALDevice::GetDefaultDevice();

  m_hRenderScreenShader.Invalidate();
  m_hDisplayScreenShader.Invalidate();
  m_hClearHistogramShader.Invalidate();
  m_hComputeHistogramShader.Invalidate();
  m_hDisplayHistogramShader.Invalidate();

  m_hHistogramQuadMeshBuffer.Invalidate();

  m_hScreenRTV.Invalidate();
  device->DestroyTexture(m_hScreenTexture);
  m_hScreenTexture.Invalidate();

  device->DestroyTexture(m_hHistogramTexture);
  m_hHistogramTexture.Invalidate();

  SUPER::BeforeHighLevelSystemsShutdown();
}

void ezComputeShaderHistogramApp::CreateHistogramQuad()
{
  m_hHistogramQuadMeshBuffer = ezResourceManager::GetExistingResource<ezMeshBufferResource>("{4BEFA142-FEDB-42D0-84DC-58223ADD8C62}");

  if (!m_hHistogramQuadMeshBuffer.IsValid())
  {
    ezVec2 pixToScreen(1.0f / m_pWindow->GetClientAreaSize().width * 0.5f, 1.0f / m_pWindow->GetClientAreaSize().height * 0.5f);
    const float borderOffsetPix = 80.0f;
    const float sizeScreen = 0.8f;

    ezGeometry geom;
    ezGeometry::GeoOptions opt;
    opt.m_Color = ezColor::Black;
    opt.m_Transform = ezMat4(ezMat3::MakeIdentity(), ezVec3(1.0f - pixToScreen.x * borderOffsetPix - sizeScreen / 2, -1.0f + pixToScreen.y * borderOffsetPix + sizeScreen / 2, 0.0f));
    geom.AddRect(ezVec2(sizeScreen, sizeScreen), 1, 1, opt);

    ezMeshBufferResourceDescriptor desc;
    desc.AddCommonStreams();
    desc.AllocateStreamsFromGeometry(geom);

    ezUInt32 t = 0;
    for (ezUInt32 p = 0; p < geom.GetPolygons().GetCount(); ++p)
    {
      for (ezUInt32 v = 0; v < geom.GetPolygons()[p].m_Vertices.GetCount() - 2; ++v)
      {
        desc.SetTriangleIndices(t, geom.GetPolygons()[p].m_Vertices[0], geom.GetPolygons()[p].m_Vertices[v + 1], geom.GetPolygons()[p].m_Vertices[v + 2]);

        ++t;
      }
    }

    m_hHistogramQuadMeshBuffer = ezResourceManager::GetOrCreateResource<ezMeshBufferResource>("{4BEFA142-FEDB-42D0-84DC-58223ADD8C62}", std::move(desc));
  }
}
#if EZ_ENABLED(EZ_SUPPORTS_DIRECTORY_WATCHER)
void ezComputeShaderHistogramApp::OnFileChanged(ezStringView sFilename, ezDirectoryWatcherAction action, ezDirectoryWatcherType type)
{
  if (action == ezDirectoryWatcherAction::Modified && type == ezDirectoryWatcherType::File)
  {
    ezLog::Info("The file {0} was modified", sFilename);
    m_bStuffChanged = true;
  }
}
#endif
EZ_APPLICATION_ENTRY_POINT(ezComputeShaderHistogramApp);
